import asyncio
import json
import logging
from typing import Any
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from pydantic import BaseModel


chatApi = APIRouter(prefix='/chat', tags=["chat"])

class Message(BaseModel):
    id: str
    content: str
    role: str = 'user'

@chatApi.post('/completion')
async def completion(msg: Message) -> Any:
    logging.info(msg)
    
    def stream_reply(data: dict) -> str:
        return f'data: {json.dumps(data)}\n\n'

    async def stream():
        content = msg.content
        tmp = ""
        for c in content:
            tmp += c
            await asyncio.sleep(0.2)
            yield stream_reply({"data": tmp})
        yield stream_reply({"data": True})

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive"
        }
    )